{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using CW-Analyzer for CPA Attack\n",
    "\n",
    "Supported setups:\n",
    "\n",
    "SCOPES:\n",
    "\n",
    "* OPENADC\n",
    "* CWNANO\n",
    "\n",
    "PLATFORMS:\n",
    "\n",
    "* CWLITEARM\n",
    "* CWLITEXMEGA\n",
    "* CWNANO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial will take you through a complete attack on a software AES implementation. The specific implementation being attacked is a well-known AES implementation written in C, which is likely to be similar to other implementations used by proprietary systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET = 'TINYAES128C'\n",
    "num_traces = 50\n",
    "CHECK_CORR = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\"\n",
    "cd ../hardware/victims/firmware/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Capturing Power Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use some helper scripts to make setup and programming easier. If you're using an XMEGA or STM (CWLITEARM) target, binaries with the correct should be setup for you:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"Helper_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = '../hardware/victims/firmware/simpleserial-aes/simpleserial-aes-{}.hex'.format(PLATFORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, before we capture our traces, we'll need to create a ChipWhipserer project, since that's what Analyzer expects for an input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project = cw.create_project(\"projects/Tutorial_B5\", overwrite = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Capturing Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below you can see the capture loop. The main body of the loop loads some new plaintext, arms the scope, sends the key and plaintext, then finally records our new trace into our project (a tuple of the trace_data, textin, textout, and key)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tnrange\n",
    "ktp = cw.ktp.Basic()\n",
    "\n",
    "for i in tnrange(num_traces, desc='Capturing traces'):\n",
    "    key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "    trace = cw.capture_trace(scope, target, text, key)\n",
    "    if trace is None:\n",
    "        continue\n",
    "    project.traces.append(trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot a trace and see what it looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import holoviews as hv\n",
    "hv.extension('bokeh')\n",
    "hv.Curve(project.waves[0]).opts(width=600, height=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you'd like, you can also save the project for later analysis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're now done with the ChipWhisperer hardware, so we should disconnect from the scope and target:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleanup the connection to the target and scope\n",
    "scope.dis()\n",
    "target.dis()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "import chipwhisperer.analyzer as cwa\n",
    "proj = cw.open_project(\"projects/Tutorial_B5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our traces, we can begin our attack! We'll start by setting up our leakage model to be the output of the SBox operation (more on that in PA_CPA_2):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leak_model = cwa.leakage_models.sbox_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see a list of available leakage models by printing `cwa.leakage_models`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(cwa.leakage_models)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then the attack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attack = cwa.cpa(proj, leak_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see some of the settings for the attack by printing the attack object. You should see that we're using all of our traces, using all of the points, and attacking all 1 byte subkeys of the full 16 byte key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(attack)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can run the attack:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = attack.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then print the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should find that the KGuess column matches the known key recorded earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results.find_key())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in proj.keys[0]:\n",
    "    print(\"0x{:02X}\".format(key))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can get the full information from the attack by calling `results.find_maximums()`, which returns:\n",
    "\n",
    "```Python\n",
    "find_maxiums() ->\n",
    "    [subkey0_data, subkey1_data, subkey2_data, ...]\n",
    "    \n",
    "subkey0_data ->\n",
    "    [guess0, guess1, guess2, ...]\n",
    "    \n",
    "guess0 ->\n",
    "    (key_guess, location_of_max, correlation)\n",
    "```\n",
    "\n",
    "For example, if you want to print the correlation of the third best guess of the 4th subkey, you would run:\n",
    "\n",
    "```python\n",
    "print(attack_results.find_maximums()[4][3][2])\n",
    "```\n",
    "\n",
    "Note the \"point location of the max\" is normally not calculated/tracked, and thus returns as a 0. Using the pandas library lets us print them nicely in a DataFrame. We have to transpose the frame to get our expected orientation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "stat_data = results.find_maximums()\n",
    "df = pd.DataFrame(stat_data).transpose()\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even better, we can use the `.style` method to customize this further. This also lets us chain formatting functions. For example, we can remove the extra 0 and clean up the data. Since we know the correct key, we can even do things like printing the key in a different colour! \n",
    "\n",
    "You can do lots of formatting thanks to the pandas library! Check out https://pandas.pydata.org/pandas-docs/stable/style.html for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = proj.keys[0]\n",
    "def format_stat(stat):\n",
    "    return str(\"{:02X}<br>{:.3f}\".format(stat[0], stat[2]))\n",
    "\n",
    "def color_corr_key(row):\n",
    "    global key\n",
    "    ret = [\"\"] * 16\n",
    "    for i,bnum in enumerate(row):\n",
    "        if bnum[0] == key[i]:\n",
    "            ret[i] = \"color: red\"\n",
    "        else:\n",
    "            ret[i] = \"\"\n",
    "    return ret\n",
    "\n",
    "df.head().style.format(format_stat).apply(color_corr_key, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see red numbers printed at the top of a table. Congratulations, you've now completed a successful CPA attack against AES!\n",
    "\n",
    "Next, we'll look at how we can use some of Analyzer's other features to improve the attack process, as well as better interpret the data we have."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reporting Intervals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we ran `attack.run()`, we processed all of the traces before getting any information back. While this works okay for shorter attacks like this, for longer ones it can helpful to get feedback during the attack. This can be done by creating a callback function and passing it to `attack.run()`. This function is called each time we pass the update interval (default 25, which is the second parameter for `attack.run()`).\n",
    "\n",
    "Let's use this to update our table every 10 traces. Most of this is just putting our existing code into the callback function. We also need use the `clear_output` function to clear the table, as well as `display()` to actually get it to show up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "import numpy as np\n",
    "        \n",
    "def stats_callback():\n",
    "    results = attack.results\n",
    "    results.set_known_key(key)\n",
    "    stat_data = results.find_maximums()\n",
    "    df = pd.DataFrame(stat_data).transpose()\n",
    "    clear_output(wait=True)\n",
    "    display(df.head().style.format(format_stat).apply(color_corr_key,axis=1))\n",
    "    \n",
    "results = attack.run(stats_callback, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A default jupyter callback is also available:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we used a reporting interval of 10 traces. Depending on the attack and what you want to learn from it, you may want to use higher or lower values: in general reporting less often is faster, but more frequent reporting can allow you to end a long attack early. More frequent reporting also increases the resolution of some plot data (which we will look at next)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Analyzer also includes a module to create plots to help you interpret the data. These act on one subkey at a time and return some data that we can plot using bokeh (or your graphing module of choice). Let's start by grabbing the class that does all the calculations:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Output Vs. Time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data = cwa.analyzer_plots(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start by looking at the Output Vs. Time module, which will allow us to plot correlation of our guesses in time. This is useful for finding exactly where the operations we're attacking are. Like in previous tutorials, we'll use bokeh to plot the data we get back.\n",
    "\n",
    "The method we're interested in is `get_plot_data(bnum)`, which returns in a list: `[xrange, correct_key, incorrect_key_data, incorrect_key_data]` for the position `bnum` passed to it. The method returns two sets of incorrect key data because one is for the key guesses below the correct one, and the other is for guesses above the correct one.\n",
    "\n",
    "We'll have a lot of points, so we'll plot as usual, but at the end decimtae the output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import holoviews as hv\n",
    "from holoviews.operation.datashader import datashade, shade, dynspread, rasterize\n",
    "from holoviews.operation import decimate\n",
    "import pandas as pd, numpy as np\n",
    "\n",
    "a = []\n",
    "b = []\n",
    "hv.extension('bokeh')\n",
    "for i in range(0, 16):\n",
    "    data = plot_data.output_vs_time(i)\n",
    "    a.append(np.array(data[1]))\n",
    "    b.append(np.array(data[2]))\n",
    "    b.append(np.array(data[3]))\n",
    "    \n",
    "pda = pd.DataFrame(a).transpose().rename(str, axis='columns')\n",
    "pdb = pd.DataFrame(b).transpose().rename(str, axis='columns')\n",
    "curve = hv.Curve(pdb['0']).options(color='green')\n",
    "for i in range(1, 16):\n",
    "    curve *= hv.Curve(pdb[str(i)]).options(color='green')\n",
    "\n",
    "for i in range(0, 16):\n",
    "    curve *= hv.Curve(pda[str(i)]).options(color='red')\n",
    "decimate(curve.opts(width=900, height=600))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see some distinctive red spikes in your plot. The largest of these is where the sbox lookup is actually happening (the smaller ones are typically other AES operations that move the sbox data around).\n",
    "\n",
    "Let's repeat this for all the subkeys. This is quite a bit more data to plot, so give it a few seconds:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This information can be useful in many ways. For example, you can probably see the first 16 spikes that make up the sbox lookup are a small portion of the total trace length. If we ever needed to rerun the attack, we could capture a much smaller number of samples and speed up analysis significantly!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PGE vs. Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next data we'll look at is a plot of partial guessing entropy (PGE) vs. the number of traces. As mentioned before, PGE is just how many spots away from the top the actual subkey is in our table of guesses. For example, if there are 7 subkey guesses that have a higher correlation than the actual subkey, the subkey has a PGE of 7.\n",
    "\n",
    "This plot is useful for seeing how many traces were needed to actually break the AES implementation. Keep in mind, however, that the resolution of the plot is determined by the reporting interval (also note that `attack_results.find_maximums()` must be called in the callback function). In our case, we have a reporting interval of 10, so we'll have a resolution of 10 traces.\n",
    "\n",
    "This method is similar to the previous plot in that it takes `bnum` as an argument and returns a list of `[xrange, PGE]`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = plot_data.pge_vs_trace(0)\n",
    "curve = hv.Curve((ret[0],ret[1]))\n",
    "for bnum in range(1, 16):\n",
    "    ret = plot_data.pge_vs_trace(bnum)\n",
    "    curve *= hv.Curve((ret[0],ret[1]))\n",
    "curve.opts(width=900, height=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see a number of lines that start off with high values, then rapidly drop off. You may notice that we broke the AES implementation without needing to use all of our traces. \n",
    "\n",
    "Even though we may have broken the AES implementation in fewer traces, we may not want to reduce how many traces we capture. Remember that, while we know the key here, for a real attack we won't and therefore must use the correlation to determine when we've broken a key. Our next plot will help us to determine how feesible capturing fewer traces is."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Correlation vs. Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last plot we'll take a look at is correlation vs the number of traces. Like with PGE vs. Traces, this plot's resolution is determined by the reporting interval (10 in our case). One again, this is a plot with a lot of data, so we'll start of by just plotting one subkey:\n",
    "\n",
    "This method returns a list of `[xrange, [data_for_kguess]]`, so we'll need to plot each guess for each subkey. Like before, we'll do the plot for the correct subkey in red and the rest in green."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, all the subkey guesses start of with large correlations, but all of them except for the correct guess quickly drop off. If you didn't know the key, at what point would you be sure that the guess with the highest correlation was actually the correct subkey?\n",
    "\n",
    "Let's continue and plot the correlations for the right guess and the next best one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = []\n",
    "b = []\n",
    "for bnum in range(0, 16):\n",
    "    data = plot_data.corr_vs_trace(bnum)\n",
    "    best = [0] * len(data[1][0])\n",
    "    for i in range(255):\n",
    "        if i == key[bnum]:\n",
    "            a.append(np.array(data[1][i]))\n",
    "        else:\n",
    "            if max(best) < max(data[1][i]): best = data[1][i]\n",
    "    b.append(np.array(data[1][i]))\n",
    "\n",
    "pda = pd.DataFrame(a).transpose().rename(str, axis='columns')\n",
    "pdb = pd.DataFrame(b).transpose().rename(str, axis='columns')\n",
    "curve = hv.Curve(pdb['0']).options(color='green')\n",
    "for i in range(1,len(pdb.columns)):\n",
    "    curve *= hv.Curve(pdb[str(i)]).options(color='green')\n",
    "    \n",
    "for i in range(len(pda.columns)):\n",
    "    curve *= hv.Curve(pda[str(i)]).options(color='red')\n",
    "            \n",
    "curve.opts(width=900, height=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like in the first plot, you should see the red lines remain high while the green ones drop off. At what point would you be sure that you've broken all the subkeys? Is it higher than when all of the PGE lines reached zero?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now have completed a successful CPA attack and learned about some on Analyzer's features for improving your attack! \n",
    "\n",
    "You can move onto more advanced tutorials, especially showing you how the actual attack works when performed manually (Tutorial B6). This tutorial also utilized tiny-AES128-C for Arm targets, which uses the same operations as the XMEGA target. A later tutorial will preform this attack on a more typical 32 bit AES implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = list(proj.keys[0])\n",
    "recv_key = [kguess[0][0] for kguess in results.find_maximums()]\n",
    "assert (key == recv_key), \"Failed to recover encryption key\\nGot: {}\\nExpected: {}\".format(recv_key, key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert (results.pge == [0]*16), \"PGE for some bytes not zero: {}\".format(results.pge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if CHECK_CORR:\n",
    "    max_corrs = [kguess[0][2] for kguess in results.find_maximums()]\n",
    "    assert (np.all([corr > 0.75 for corr in max_corrs])), \"Low correlation in attack (corr <= 0.75): {}\".format(max_corrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
